cmake_minimum_required(VERSION 3.17)
project(CornetNet_Lite_TRT)

set(CMAKE_CXX_STANDARD 14)

find_package(CUDA REQUIRED)

find_package(OpenCV REQUIRED)
if (NOT OpenCV_FOUND)
    message(FATAL_ERROR "OpenCV library not found")
endif()

# CUDA PATH
set(CUDA_INC_PATH "/usr/local/cuda/include")
set(CUDA_LIB_PATH "/usr/local/cuda/lib64")

# TENSORRT PATH
set(TRT_INC_PATH "/usr/local/TensorRT/include")
set(TRT_LIB_PATH "/usr/local/TensorRT/lib")

# Compiler Options
if (NOT CMAKE_BUILD_TYPE)
    set(CMAKE_BUILD_TYPE "Release")
endif()

if (CMAKE_BUILD_TYPE STREQUAL Release)
    set(CMAKE_CXX_FLAGS "-O2 -ffast-math")
endif()

set(CMAKE_CXX_FLAGS "$ENV{CXXFLAGS} -Wno-deprecated-declarations ${CMAKE_CXX_FLAGS}")
set(CUDA_NVCC_FLAGS "--std=c++14
                     -gencode arch=compute_60,code=sm_60
                     -gencode arch=compute_61,code=sm_61
                     -gencode arch=compute_70,code=sm_70
                     -gencode arch=compute_75,code=sm_75")

include_directories(
        ${CUDA_INC_PATH}
        ${TRT_INC_PATH}
        ${OpenCV_INCLUDE_DIRS}
        ${CMAKE_CURRENT_SOURCE_DIR}/include/
)
link_directories(
        ${CUDA_LIB_PATH}
        ${TRT_LIB_PATH}
)

set(PLUGIN_SRCS
        plugins/corner_pool_plugin.h
        plugins/left_pool_plugin.h
        plugins/bottom_pool_plugin.h
        plugins/right_pool_plugin.h
        plugins/top_pool_plugin.h
        plugins/corner_pool.cu
        )

cuda_add_executable(CornetNetLiteTrtExec ${PLUGIN_SRCS} trtexec.cpp)
target_link_libraries(CornetNetLiteTrtExec
        cudart
        nvinfer
        nvparsers
        nvinfer_plugin
        nvonnxparser
        ${OpenCV_LIBS})

cuda_add_executable(TrtInference ${PLUGIN_SRCS} trt_inference.cpp)
target_link_libraries(TrtInference
        cudart
        nvinfer
        nvparsers
        nvinfer_plugin
        nvonnxparser
        ${OpenCV_LIBS})
